import nibabel as nib
import matplotlib.pyplot as plt
import numpy as np

def view_nifti(nifti_path, seg_map=None):  # 添加一个参数用于传递分割地图数据
    # 读取.nii.gz文件
    img = nib.load(nifti_path)
    # 获取图像数据
    data = img.get_fdata()

    # 显示图像的多个切片
    num_slices = data.shape[-1]
    num_rows = num_slices // 5  # 将多个切片分成5行
    fig, axes = plt.subplots(num_rows, 5, figsize=(15, 3 * num_rows))

    for i in range(num_slices):
        ax = axes[i // 5, i % 5]  # 计算当前切片应该在哪个子图上显示
        # 将数据归一化到 [0, 1] 范围内
        # 由于imshow的cmap需要数据在[0, 1]范围内，所以这里进行归一化
        # 注意：归一化的方法可能需要根据数据的实际情况进行调整
        gray_image_normalized = (data[:, :, i] - np.min(data[:, :, i])) / (np.max(data[:, :, i]) - np.min(data[:, :, i]))

        # 使用'viridis'颜色映射显示彩色图像
        ax.imshow(gray_image_normalized, cmap='viridis')  # 显示彩色图像

        # 如果有分割地图，可以使用不同的颜色表示不同的标签
        if seg_map is not None:
            ax.imshow(seg_map[:, :, i], alpha=0.5)  # 将分割地图叠加在彩色图像上，使用透明度来区分

        ax.axis('off')  # 不显示坐标轴
        ax.set_title(f"Slice {i}")  # 设置子图标题

    plt.show()

# 调用函数查看图像
# view_nifti(r'/root/autodl-tmp/BrainTumorSegmentation1/predictions/UNet/000_pred.nii.gz')
view_nifti(r'D:\PyCharmCode\BrainTumorSegmentation1\predictions\UNet1\000_gt.nii.gz')
#view_nifti(r'D:\PyCharmCode\BrainTumorSegmentation1\predictions\UNet1\000_img.nii.gz')
#view_nifti(r'D:\PyCharmCode\BrainTumorSegmentation1\predictions\UNet1\000_pred.nii.gz')
